﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using TorchSharp;

namespace NgpSharp.RayMaching
{
    /// <summary>
    /// 
    /// </summary>
    public class NearFarAABBFunction : torch.autograd.MultiTensorFunction<NearFarAABBFunction>
    {
        /// <summary>
        /// 
        /// </summary>
        public override string Name { get => "near_far_from_aabb"; }
        /// <summary>
        /// 
        /// </summary>
        /// <param name="ctx"></param>
        /// <param name="grad_outputs"></param>
        /// <returns></returns>
        public override List<torch.Tensor> backward(torch.autograd.AutogradContext ctx, List<torch.Tensor> grad_outputs)
        {
            foreach (var item in grad_outputs) 
            {
                item.Dispose();
            }
            return [null, null, null, null];
        }
        /// <summary>
        /// 
        /// </summary>
        /// <param name="ctx"></param>
        /// <param name="vars"></param>
        /// <returns></returns>
        public override List<torch.Tensor> forward(torch.autograd.AutogradContext ctx, params object[] vars)
        {
            using var c = torch.NewDisposeScope();
            torch.Tensor rays_o = (torch.Tensor)vars[0]; 
            torch.Tensor rays_d = (torch.Tensor)vars[1];
            torch.Tensor aabb = (torch.Tensor)vars[2];
            float min_near = (float)vars[3];
            if(!rays_o.is_cuda) rays_o = rays_o.cuda();
            if(!rays_d.is_cuda) rays_d = rays_d.cuda();

            rays_o = rays_o.contiguous().view(-1, 3);
            rays_d = rays_d.contiguous().view(-1, 3);

            long N = rays_o.shape[0];

            torch.Tensor nears = torch.empty(N, dtype: rays_o.dtype, device: rays_o.device);
            torch.Tensor fars = torch.empty(N, dtype: rays_d.dtype, device: rays_d.device);

            RayMachingLoader.near_far_from_aabb(rays_o.Handle, rays_d.Handle, aabb.Handle, (uint)N, min_near, nears.Handle, fars.Handle);

            return [nears.MoveToOuterDisposeScope(), fars.MoveToOuterDisposeScope()];
        }
    }
}
